#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2019/5/15 11:12
# @Author  : CHEN Wang
# @Site    : 
# @File    : my_excel.py
# @Software: PyCharm

"""
脚本说明: 对Excel的读写等操作
"""

import os
import csv
import pandas as pd
from openpyxl import load_workbook
from openpyxl.styles import Font, colors, Alignment
from quant_researcher.quant.project_tool.logger.my_logger import LOG
from quant_researcher.quant.project_tool.db_operator import my_mysql
from quant_researcher.quant.project_tool import conf, hammer

# BOLD_ITALIC_24_FONT = Font(name='等线', size=24, italic=True, color=colors.RED, bold=True)
# 颜色默认为红色
BOLD_ITALIC_24_FONT = Font(name='等线', size=24, italic=True, color='FF0000', bold=True)


def read_xl_and_write_db(file_path, primary_key, table, conn, **kwargs):
    """
    读取Excel并存入数据库

    :param file_path: Excel文件路径与文件名
    :param list primary_key:
    :param table: 数据库表名
    :param conn: 数据库连接
    :param kwargs: 读取read_excel需要传入的参数
            - sheetname
            - header
            - ...
    :return: None
    """

    sheetname = kwargs.pop('sheet_name', 0)
    df = pd.read_excel(file_path, sheet_name=sheetname)
    hammer.slim(df, show_index=1, a_f=1)
    my_mysql.on_duplicate_key_update(a_lst=df, primary_key=primary_key, conn=conn, table=table,
                                     replace_nan=True)

    return None


def df_2_excel(df, res_dir=None, file_name='output', **kwargs):
    """
    输出DataFrame到excel, 为什么有？路径组合和输出提示的代码复用。2020年3月22日16:11:01

    :param pd.DataFrame df: 需要输出的DataFrame
    :param res_dir: 文件输出位置
    :param file_name: 文件名称,不用加后缀，自动加
    :param kwargs: to_excel需要传入的参数
           - index，bool，是否将DF的index一并写入文件，默认：False
    :return:
    """
    sheetname = kwargs.pop('sheet_name', 'Sheet1')  # 默认创建工作表名为Sheet1
    if_index = kwargs.pop('index', False)
    if res_dir is None:
        res_dir = conf.BASE_DIR
        # res_dir = '/Users/Damien/Desktop'
    if os.path.exists(res_dir):
        LOG.warning(f'预计的结果存储目录已存在：{res_dir}')
    else:
        LOG.warning(f'预计的结果存储目录不存在，即将创建：{res_dir}')
        os.mkdir(res_dir)
        LOG.warning(f'已创建目录：{res_dir}')
    res_path = os.path.join(res_dir, f'{file_name}.xlsx')
    print('🚂🚃: 将要把DF写入文件(kwargs: %s)：%s' % (kwargs, file_name))
    df.to_excel(res_path, encoding='gbk', sheet_name=sheetname, index=if_index, freeze_panes=(0, 0))
    print('🚅🛂: 写入成功！')


def df_2_sheet_in_excel(df, file_name=None, sheet_name=None, **kwargs):
    """
    df 写到Excel中的某个sheet里面，不影响其他的sheet

    :param pd.DataFrame df: 需要输出的DataFrame
    :param str file_name: 文件地址+文件名, 必须输入， 如os.path.join(f"E:\指标数据\ohlcv", 'test.xlsx')
    :param sheet_name: 工作表名，必须指定
    :param kwargs: 读取to_excel需要传入的参数
        - index，bool，是否将DF的index一并写入文件，默认：False
    :return:
    """

    if_index = kwargs.pop('index', False)
    if not os.path.exists(file_name):
        LOG.error(f'目标表{file_name}不存在')

    # 将原表中的内容读取出来
    writer_book = load_workbook(f'{file_name}.xlsx')
    # 将原表中的内容保存到writer
    writer = pd.ExcelWriter(f'{file_name}.xlsx', engine='openpyxl')
    writer.book = writer_book
    idx = writer_book.sheetnames.index(sheet_name)  # 找到sheet_name的索引
    writer_book.remove(writer_book.worksheets[idx])  # 删除sheet_name工作表中的原数据
    writer_book.create_sheet(sheet_name, idx)  # 重新创建空白的sheet_name工作表
    writer.sheets = dict((i.title, i) for i in writer_book.worksheets)  # 保存其他用不到的sheet页面
    # 将新数据data写入sheet_name工作表中
    df.to_excel(excel_writer=writer, sheet_name=sheet_name, encoding='gbk', index=if_index, freeze_panes=(0, 0))
    writer.save()
    writer.close()
    print(f'🚅🛂: 数据成功更新到{writer_book.worksheets[idx]}')


def df_list_2_excel_sheets(df_list, res_dir=None, file_name='output', sheet_name_list=None, **kwargs):
    """
    df_list 写到同一个Excel不同sheet里面； 注意如果Excel中原来包含多个sheet, 原有数据会被覆盖，无法只更新某个sheet

    :param df_list: list, df列表
    :param res_dir: 文件输出位置
    :param file_name: 文件名称,不用加后缀，自动加
    :param sheet_name_list: sheet_name名称列表
    :param kwargs: 读取to_excel需要传入的参数
        - index，bool，是否将DF的index一并写入文件，默认：False
    :return:
    """
    if_index = kwargs.pop('index', False)
    if res_dir is None:
        res_dir = conf.BASE_DIR
    if os.path.exists(res_dir):
        LOG.warning(f'预计的结果存储目录已存在：{res_dir}')
    else:
        LOG.warning(f'预计的结果存储目录不存在，即将创建：{res_dir}')
        os.mkdir(res_dir)
        LOG.warning(f'已创建目录：{res_dir}')

    target_path = os.path.join(res_dir, f'{file_name}.xlsx')
    writer = pd.ExcelWriter(target_path, engine='openpyxl')

    if len(df_list) != len(sheet_name_list):
        LOG.warning(f'传入的df_list长度与sheet_name_list不一致')

    print(f'🚂🚃: 将要把DF写入文件(kwargs: {kwargs})：{res_dir}->{file_name}')
    for df, sheet_name in zip(df_list, sheet_name_list):
        df.to_excel(excel_writer=writer, sheet_name=sheet_name, encoding='gbk', index=if_index,
                    freeze_panes=(0, 0))
    writer.save()
    writer.close()
    print('🚅🛂: 写入成功！')


def QA_util_save_csv(data, name, column=None, location=None):
    # 重写了一下保存的模式
    # 增加了对于可迭代对象的判断 2017/8/10
    """
    explanation:
        将数据保存为csv

    params:
        * data->
            含义: 需要被保存的数据
            类型: List
            参数支持: []
        * name:
            含义: 要保存的名字
            类型: str
            参数支持: []
        * column:
            含义: 行的名称(可选)
            类型: str
            参数支持: [None]
        * location:
            含义: 保存位置(可选)
            类型: str
            参数支持: []
    """
    assert isinstance(data, list)
    if location is None:
        path = './' + str(name) + '.csv'
    else:
        path = location + str(name) + '.csv'
    with open(path, 'w', newline='') as f:
        csvwriter = csv.writer(f)
        if column is None:
            pass
        else:
            csvwriter.writerow(column)

        for item in data:

            if isinstance(item, list):
                csvwriter.writerow(item)
            else:
                csvwriter.writerow([item])


if __name__ == '__main__':
    # 测试 df_list_2_excel_sheet
    passive_fund_list = pd.DataFrame([['2020-03-31', '110022', 0.5], ['2020-03-31', '001410', 0.5],
                                      ['2020-05-03', '110022', 0.6], ['2020-05-03', '001410', 0.4]],
                                     columns=['tradedate', 'code', 'weights'])
    active_stock_fund_list = pd.DataFrame([['2020-03-31', '110022', 0.5], ['2020-03-31', '001410', 0.5],
                                           ['2020-05-03', '110022', 0.6], ['2020-05-03', '001410', 0.4]],
                                          columns=['tradedate', 'code', 'weights'])

    fund_list = [passive_fund_list, active_stock_fund_list]
    sheet_name_list = ['被动基金榜单', '主动基金榜单']

    file_name = '基金榜单'
    desktop_dir = os.path.join(os.path.expanduser("~"), 'Desktop')

    df_list_2_excel_sheets(df_list=fund_list, res_dir=desktop_dir, file_name=file_name,
                           sheet_name_list=sheet_name_list)
    # 测试 QA_util_save_csv
    QA_util_save_csv(['a', 'v', 2, 3], 'test')
    QA_util_save_csv([['a', 'v', 2, 3]], 'test2')
